"""
Routing Chain Service Demo

This service demonstrates the use of LangChain's routing chain pattern,
which can automatically route queries to different processing chains based on content.
"""

from typing import Dict, List, Any, Literal
from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
from langchain_core.output_parsers import StrOutputParser, JsonOutputParser
from langchain.chains import LLMChain
from langchain.chains.router import MultiPromptChain
from langchain.chains.router.llm_router import LLMRouterChain, RouterOutputParser
from langchain.chains.router.multi_prompt_prompt import MULTI_PROMPT_ROUTER_TEMPLATE
from langchain_openai import ChatOpenAI

from app.services.zodiac_service import get_zodiac_sign, get_all_zodiac_signs, get_zodiac_compatibility
from app.services.system_agent import handle_system_query
from app.core.config import settings


class RoutingChainService:
    """
    Routing Chain Service for demonstrating different routing patterns
    """

    def __init__(self, model_name: str = "deepseek-chat", temperature: float = 0.7):
        self.model_name = model_name
        self.temperature = temperature
        self.llm = ChatOpenAI(
            model=model_name,
            temperature=temperature,
            max_tokens=65536,
            api_key=settings.openai_api_key,
            base_url=settings.openai_base_url
        )
        self.parser = StrOutputParser()

    def create_zodiac_chain(self):
        """Create chain for zodiac-related queries"""
        prompt = ChatPromptTemplate.from_messages([
            ("system", """你是一个生肖专家，专门回答关于生肖的问题。
请根据用户的查询提供详细的生肖信息，包括特质、幸运数字、颜色等。
如果需要具体的年月信息来查询生肖，请询问用户提供。"""),
            ("human", "{query}")
        ])

        def zodiac_func(query):
            # Try to extract year from query for zodiac calculation
            import re
            year_match = re.search(r'\b(19|20)\d{2}\b', query)
            if year_match:
                year = int(year_match.group())
                zodiac_info = get_zodiac_sign(year)
                if "error" not in zodiac_info:
                    return f"""
根据查询"{query}"，我为您提供以下生肖信息：

🐉 生肖: {zodiac_info['zodiac_sign']} ({zodiac_info['zodiac_english']})
🐲 生肖年: {zodiac_info['zodiac_year']}年
⭐ 特质: {zodiac_info['attributes']['traits']}
🎰 幸运数字: {', '.join(map(str, zodiac_info['attributes']['lucky_numbers']))}
🎨 幸运颜色: {', '.join(zodiac_info['attributes']['lucky_colors'])}
                    """.strip()

            # If no year found or query is general, use LLM
            return (prompt | self.llm | self.parser).invoke({"query": query})

        return zodiac_func

    def create_system_chain(self):
        """Create chain for system-related queries"""
        prompt = ChatPromptTemplate.from_messages([
            ("system", """你是一个系统监控助手，专门回答关于系统状态、性能、会话统计等问题。
你可以查询和提供以下信息：
- 系统负载（CPU、内存、磁盘使用情况）
- 会话统计（总会话数、最近会话、消息数量）
- 数据库状态和性能
- 服务器运行时间和健康状态

请根据用户的具体查询提供准确的系统信息。"""),
            ("human", "{query}")
        ])

        def system_func(query):
            # Use the system agent to handle the query
            try:
                result = handle_system_query(query)
                return result
            except Exception as e:
                # Fallback to LLM if system agent fails
                return (prompt | self.llm | self.parser).invoke({"query": query})

        return system_func

    def create_general_chain(self):
        """Create chain for general conversation"""
        prompt = ChatPromptTemplate.from_messages([
            ("system", """你是一个智能助手，可以回答各种问题。
你能够：
- 进行日常对话
- 回答常识性问题
- 提供有用的信息和建议
- 进行数学计算
- 解释概念和定义

请用友好、自然的方式与用户交流。"""),
            ("human", "{query}")
        ])

        return prompt | self.llm | self.parser

    def create_calculation_chain(self):
        """Create chain for mathematical calculations"""
        prompt = ChatPromptTemplate.from_messages([
            ("system", """你是一个数学专家，专门处理各种数学计算。
你可以：
- 进行基础算术运算
- 解决代数问题
- 计算百分比和比率
- 处理简单的几何问题
- 进行单位转换

请确保计算结果准确，并给出清晰的解答过程。"""),
            ("human", "{query}")
        ])

        def calc_func(query):
            # First try to evaluate simple expressions
            import re
            import ast

            # Extract mathematical expressions
            expr_match = re.search(r'([\d+\-*/().\s]+)', query)
            if expr_match:
                expr = expr_match.group(1).strip()
                try:
                    # Safe evaluation of mathematical expressions
                    result = eval(expr, {'__builtins__': None}, {})
                    return f"计算结果: {expr} = {result}"
                except:
                    pass

            # Fallback to LLM for complex calculations
            return (prompt | self.llm | self.parser).invoke({"query": query})

        return calc_func

    def create_routing_chain_multi_prompt(self):
        """
        Create a routing chain using MultiPromptChain (LangChain's built-in router)
        """
        # Define destination chains
        destination_chains = {
            "zodiac": self.create_zodiac_chain(),
            "system": self.create_system_chain(),
            "calculation": self.create_calculation_chain(),
            "general": self.create_general_chain()
        }

        # Define prompt templates for each destination
        prompt_infos = [
            {
                "name": "zodiac",
                "description": "适合处理关于生肖查询、生肖配对、农历相关的问题",
                "prompt_template": "用户询问关于生肖的问题: {query}"
            },
            {
                "name": "system",
                "description": "适合处理系统监控、服务器负载、会话统计、数据库状态等系统相关问题",
                "prompt_template": "用户询问关于系统状态的问题: {query}"
            },
            {
                "name": "calculation",
                "description": "适合处理数学计算、算术运算、百分比计算等数学相关问题",
                "prompt_template": "用户询问数学计算问题: {query}"
            },
            {
                "name": "general",
                "description": "适合处理日常对话、常识性问题、一般咨询等非专业领域的问题",
                "prompt_template": "用户的一般性问题: {query}"
            }
        ]

        # Create the multi-prompt chain
        router_chain = MultiPromptChain.from_prompts(
            self.llm,
            prompt_infos=prompt_infos,
            destination_chains=destination_chains,
            default_chain=destination_chains["general"]
        )

        return router_chain

    def create_custom_routing_chain(self):
        """
        Create a custom routing chain with more intelligent routing logic
        """
        # Define the router prompt
        router_prompt = PromptTemplate.from_template("""
你是一个智能路由器，需要根据用户的问题将其路由到最合适的处理链。

可用的处理链:
1. zodiac - 处理生肖查询、生肖配对、农历相关问题
2. system - 处理系统监控、服务器负载、会话统计、数据库状态
3. calculation - 处理数学计算、算术运算、百分比计算
4. general - 处理日常对话、常识性问题、一般咨询

请分析以下问题，并选择最合适的处理链。只返回链的名称，不要其他内容。

问题: {query}

选择的链:""")

        # Define destination chains
        chains = {
            "zodiac": self.create_zodiac_chain(),
            "system": self.create_system_chain(),
            "calculation": self.create_calculation_chain(),
            "general": self.create_general_chain()
        }

        def route_and_execute(query: str) -> str:
            # First, determine which chain to use
            router_input = router_prompt.format(query=query)
            router_response = self.llm.invoke(router_input)
            chain_name = router_response.content.strip()

            # Validate chain name
            if chain_name not in chains:
                chain_name = "general"  # Default fallback

            # Execute the selected chain
            selected_chain = chains[chain_name]

            try:
                if callable(selected_chain):
                    result = selected_chain(query)
                else:
                    result = selected_chain.invoke({"query": query})

                return f"[路由到: {chain_name}]\n{result}"
            except Exception as e:
                # Fallback to general chain if selected chain fails
                fallback_result = chains["general"].invoke({"query": query})
                return f"[路由到: {chain_name} 失败，回退到 general]\n{fallback_result}"

        return route_and_execute

    def create_semantic_routing_chain(self):
        """
        Create a semantic routing chain that uses keyword matching and intent analysis
        """
        def semantic_router(query: str) -> str:
            query_lower = query.lower()

            # Define routing rules based on keywords and intent
            zodiac_keywords = ["生肖", "属相", "鼠", "牛", "虎", "兔", "龙", "蛇", "马", "羊", "猴", "鸡", "狗", "猪", "本命年", "配对"]
            system_keywords = ["系统", "服务器", "负载", "会话", "统计", "数据库", "性能", "内存", "cpu", "磁盘", "监控"]
            calculation_keywords = ["计算", "等于", "加", "减", "乘", "除", "百分比", "多少", "数学"]

            # Route based on keyword matching
            if any(keyword in query_lower for keyword in zodiac_keywords):
                chain_name = "zodiac"
            elif any(keyword in query_lower for keyword in system_keywords):
                chain_name = "system"
            elif any(keyword in query_lower for keyword in calculation_keywords):
                chain_name = "calculation"
            else:
                chain_name = "general"

            # Get the corresponding chain
            chains = {
                "zodiac": self.create_zodiac_chain(),
                "system": self.create_system_chain(),
                "calculation": self.create_calculation_chain(),
                "general": self.create_general_chain()
            }

            selected_chain = chains[chain_name]

            try:
                if callable(selected_chain):
                    result = selected_chain(query)
                else:
                    result = selected_chain.invoke({"query": query})

                return f"[语义路由到: {chain_name}]\n{result}"
            except Exception as e:
                # Fallback
                fallback_result = chains["general"].invoke({"query": query})
                return f"[语义路由失败，使用general]\n{fallback_result}"

        return semantic_router

    def demo_routing_patterns(self, query: str) -> Dict[str, str]:
        """
        Demonstrate different routing patterns with the same query
        """
        results = {}

        # Test MultiPromptChain routing
        try:
            multi_prompt_chain = self.create_routing_chain_multi_prompt()
            results["multi_prompt_chain"] = multi_prompt_chain.run(query)
        except Exception as e:
            results["multi_prompt_chain"] = f"Error: {str(e)}"

        # Test custom routing
        try:
            custom_chain = self.create_custom_routing_chain()
            results["custom_routing_chain"] = custom_chain(query)
        except Exception as e:
            results["custom_routing_chain"] = f"Error: {str(e)}"

        # Test semantic routing
        try:
            semantic_chain = self.create_semantic_routing_chain()
            results["semantic_routing_chain"] = semantic_chain(query)
        except Exception as e:
            results["semantic_routing_chain"] = f"Error: {str(e)}"

        return results


# Global service instance
routing_chain_service = RoutingChainService()


def demo_routing_patterns(query: str) -> Dict[str, str]:
    """
    Convenience function to demo routing patterns
    """
    return routing_chain_service.demo_routing_patterns(query)